# -*- coding: utf-8 -*-
# @Time    : 2024/9/9 19:59
# @Author  : Z.J. Zhang
# @Email   : zijingzhang@mail.ustc.edu.cn
# @File    : grating.py
# @Software: PyCharm

import matplotlib
import numpy
import shapely

import mot.field.laser

matplotlib.use('tkagg')
import matplotlib.pyplot as plt
from myutils.geometrical import Plane
from mot.field.laser import NonDiffractionLaserBeam, LaserInfo
from mot.field.base import RectGrid3D


# class NeedFields(enum.Enum):
#     I_incident = 0


class Sector:
    def __init__(self, center: numpy.ndarray, rout, rin,
                 start_angle, end_angle):
        """

        :param center: shape (2,)
        :param rout:
        :param rin:
        :param start_angle:
        :param end_angle:
        """
        self.center = center
        self.rout = rout
        self.rin = rin

        self.start_angle = start_angle
        self.end_angle = end_angle
        self.shape: shapely.Polygon = self.__shape()

    def __shape(self, ):
        def line(angle):
            return shapely.LineString(
                [self.center + self.rin * numpy.array([numpy.cos(angle), numpy.sin(angle)]),
                 self.center + self.rout * numpy.array([numpy.cos(angle), numpy.sin(angle)])])

        thetas = numpy.linspace(self.start_angle, self.end_angle, 100)

        def arc(r):
            return shapely.LineString(
                [
                    line_start.coords[-1],
                    *(r * numpy.array([numpy.cos(thetas), numpy.sin(thetas)]).T)
                ]
            )

        line_start = line(self.start_angle)
        arc_out = arc(self.rout)
        arc_in = arc(self.rin)
        line_end = line(self.end_angle)
        shape = shapely.Polygon(shapely.LineString([
            *numpy.array(line_start.coords),
            *numpy.array(arc_out.coords),
            *numpy.array(line_end.coords[::-1]),
            *numpy.array(arc_in.coords[::-1]),
        ])
        )

        # plt.figure()
        # plt.plot(*shape.xy,'.-')

        return shape


class GratingSection:
    def __init__(self, shape: shapely.Polygon, grating_plane: Plane,
                 # incident_k:numpy.ndarray,
                 # diffracted_ks:numpy.ndarray
                 # name: str = None
                 ):
        """

        :param shape:
        :param plane:
        :param diffracted_ks: shape (N_diffraction, 3) 每行表示对应衍射级的波矢
        """
        self.shape = shape
        self.grating_plane = grating_plane
        # self.incident_k = incident_k
        # self.diffracted_ks = diffracted_ks
        # if not name: name = "GratingSection(%s)" % get_unique_str()
        # self.name = name
        # self.intensity_field_name = "I_" + self.name

    def cache_intensity_field(self, grid: RectGrid3D, incident_laser: NonDiffractionLaserBeam,
                              info_of_diffracted_laser: LaserInfo,
                              efficiency: float
                              ):
        # assert NeedFields.I_incident.name in grid.parameter_names

        coords_in_grating_plane = self.grating_plane.intersection_with_parallel_lines(
            grid.all_points(),
            info_of_diffracted_laser.k)
        outside_grating_section = numpy.vectorize(
            # TODO: too slow
            # TODO: 目前仅支持法向量平行于WCS z-axis的光栅
            (lambda coord: self.shape.contains(shapely.Point(coord))),
            signature="(3)-> ()")(coords_in_grating_plane).astype(float)

        Idata = incident_laser.get_intensity_for_multiple_points(
            coords_in_grating_plane[:, :, :, 0],
            coords_in_grating_plane[:, :, :, 1],
            coords_in_grating_plane[:, :, :, 2]
        ) * outside_grating_section * efficiency
        grid.set_parameter_data({info_of_diffracted_laser.field_name_I: Idata})
        return grid


from mot.field.laser import LaserFieldGenerator


class LaserDiffractedByGratingSection(LaserFieldGenerator):
    def __init__(self, grating_section: GratingSection,
                 incident_laser: NonDiffractionLaserBeam,
                 diffracted_laser_info: LaserInfo,
                 efficiency: float):
        super(LaserDiffractedByGratingSection, self).__init__(diffracted_laser_info, )
        self.grating_section = grating_section
        self.incident_laser = incident_laser
        self.efficiency = efficiency

    def cache(self, grid: RectGrid3D):
        self.grating_section.cache_intensity_field(grid, self.incident_laser, self.laser_info, self.efficiency)
        # return grid


if __name__ == '__main__':
    plt.ion()
    sector = Sector(numpy.array([0, 0]), 5e-3, 1e-3, numpy.deg2rad(-60), numpy.deg2rad(60))
    plt.plot(*sector.shape.exterior.xy)
    plt.gca().set_aspect('equal')

    grid = RectGrid3D(numpy.array((*numpy.linspace(-10e-3, -5e-3, 5),
                                   *numpy.linspace(-4.5e-3, 4.5e-3, 40),
                                   *numpy.linspace(5e-3, 10e-3, 5)
                                   )),
                      numpy.array((*numpy.linspace(-15e-3, -5e-3, 6),
                                   *numpy.linspace(-4.5e-3, 4.5e-3, 40),
                                   *numpy.linspace(5e-3, 10e-3, 5)
                                   )),
                      numpy.linspace(-10e-3, 10e-3, 22),
                      [
                          # "Bx", "By", "Bz"
                      ]
                      )
    # grid.add_new_parameter_names([f.name for f in NeedFields])
    gs = GratingSection(sector.shape, Plane(numpy.array([0, 0, grid.zs.min()]), numpy.array([0, 0, 1])))
    aperture = 8e-3
    incident_laser = NonDiffractionLaserBeam(
        LaserInfo(numpy.array([0, 0, -1]), 1,),
        Plane(numpy.array([0, 0, grid.zs.max()]), numpy.array([0, 0, 1]), ),
        lambda x, y, z: 2. if x ** 2 + y ** 2 < aperture ** 2 else 0.,
    )
    laser_diffracted_by_grating_section = LaserDiffractedByGratingSection(gs, incident_laser,
                                                                          LaserInfo(numpy.array([-1, 0, 2]), 1,
                                                                                    "grating1"), 0.4)
    laser_diffracted_by_grating_section.cache(grid)
    incident_laser.cache(grid)

    from myutils.test_utils import timer


    @timer
    def get_Idata():
        return grid.get(X, 0, Z, [laser_diffracted_by_grating_section.laser_info.field_name_I])[:, :, 0]


    X, Z = numpy.meshgrid(grid.xs, grid.zs, indexing='ij')
    plt.figure()
    plt.contourf(X / 1e-3, Z / 1e-3, get_Idata())
    plt.contourf(X / 1e-3, Z / 1e-3, grid.get(X, 0, Z, [laser_diffracted_by_grating_section.laser_info.field_name_I])[:, :, 0])
    plt.scatter(X / 1e-3, Z / 1e-3, s=0.1)
    plt.gca().set_aspect('equal')

    _X, _Y = numpy.meshgrid(grid.xs, grid.ys, indexing='ij')
    plt.figure()
    plt.contourf(_X / 1e-3, _Y / 1e-3,
                 grid.get(_X, _Y, 0, [laser_diffracted_by_grating_section.laser_info.field_name_I])[:, :, 0])
    plt.scatter(_X / 1e-3, _Y / 1e-3, s=0.1)
    plt.gca().set_aspect('equal')

    grid.save()
